from peft import LoraConfig, TaskType
from transformers import (
    TrainingArguments,
)
class Config:

    question_dict =  {
        'q1': '图中施工处于路边、农田、森林等哪类场景',
        'q2': '“施工围挡”包括金属围栏、尼龙网、雪糕筒、警戒线、金属板等，判断图中是否有“施工围挡”，如果有，具体说明有哪类围栏。',
        'q3': '图中有没有吊车？有哪些车辆。',
        'q4': '图中有没有吊物？吊物指“吊起来的重物”。',
        'q5': '工作人员站在被“吊起来的重物”下方或吊臂下方时，属于违章；作业人员扶着吊物时，属于违章；没有吊物时，工作人员站在吊臂下属于合规；其他情况为合规。',
        'q6': '在路边施工时，需要在施工场所周围设置围栏，否则属于违章行为。判断图片内容属于违章还是合规。',
        'q7': '在坑洞附近施工时，应设置围栏。否则属于违章行为。判断图片内容属于违章还是合规。 ',
        'q8': '施工现场有吊车时，若附近没有设置“施工围挡”，则属于违章行为，否则为合规。警戒线绑在吊车的脚撑上时，也属于违章。判断图中内容属于违章还是合规。',
    }

    model_path = './Qwen/Qwen2-VL-2B-Instruct/'
    dataset_path = [r'C:\sym\src\multimodal\qw_scripts\datasets\Integration.json']
    train_val_ratio = 0.8
    pic_path = [r'C:\sym\datasets\multimodal\NK66V1-001（20241120-围栏&吊车多模态数据标注任务-第1批提交247张）\NK66V1-001（20241120-围栏&吊车多模态数据标注任务-第1批提交247张）\images\Integration_img']
    random_select_question_prob = {
        'q1': 0.2,
        'q2': 0.2,
        'q3': 0.2,
        'q4': 0.2,
        'q5': 0.99,
        'q6': 0.99,
        'q7': 0.99,
        'q8': 0.99,

    }

    lora_cofig  = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        inference_mode=False,  # 训练模式
        r=64,  # Lora 秩
        lora_alpha=16,  # Lora alaph，具体作用参见 Lora 原理
        lora_dropout=0.05,  # Dropout 比例
        bias="none",
    )
 
    output_dir = './output/Qwen2-VL-2B'
    train_args = TrainingArguments(
        output_dir=output_dir,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=2,
        logging_steps=10,
        num_train_epochs=2,
        save_steps=100,
        learning_rate=1e-4,
        save_on_each_node=True,
        gradient_checkpointing=True,
        report_to="none",
        )
    
class ImageAugmentation:
    def __call__(self, *args, **kwds):
        pass
    
 